#coding: utf-8

import maya.cmds as cmds  # pylint: disable=E0401
import maya.mel as mel  # pylint: disable=E0401
import functools
import tempfile
import logging
import logging.handlers
from PSOFT.PencilSmoothPreviewVisibility import VisibleJudgment


def setup_logger():
#    cmds.optionVar(iv=('Pencil.SmoothDebug', 1))
#    cmds.optionVar(remove='Pencil.SmoothDebug')
    if not cmds.optionVar(q='Pencil.SmoothDebug'):
        class _Dummy(object):
            def info(self, *args):
                pass
        return _Dummy()

    logger = logging.getLogger(__name__)
    log_level = logging.INFO
    logger.setLevel(log_level)

    if not logger.handlers:
        logfile = tempfile.gettempdir()
        logfile += '\\smooth_batch.log' if cmds.about(q=True, batch=True) else '\\smooth.log'
        fileH = logging.handlers.RotatingFileHandler(logfile, maxBytes=0x9FFFFF, backupCount=2)
        fileH.setLevel(log_level)
        # fmt = logging.Formatter(
        # '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fmt = logging.Formatter(
            '%(asctime)s - %(levelname)s - %(module)s - ln.%(lineno)d - %(funcName)s - %(message)s')
        fileH.setFormatter(fmt)
        logger.addHandler(fileH)
    return logger

plogger = setup_logger()



#プログレスバー（Mayaメインのプログレスバーを操作します）
class MainProgress(object):
    def __init__(self, max_count=1, interruptable=False, status=None):
        try:
            bar = None
            if not cmds.about(q=True, batch=True):
                bar = mel.eval('$tmp = $gMainProgressBar')
        except RuntimeError:
            bar = None
        self.max_count = max_count if max_count > 0 else 1
        self.count = 0
        self.status = status if status is not None else u'Progress ...'
        self.bar = cmds.progressBar(bar
                                    , edit = True
                                    , beginProgress = True
                                    , isInterruptable = interruptable
                                    , status = self.status
                                    , minValue = self.count
                                    , maxValue = self.max_count
                                    )

    def step(self, stts = None):
        if stts is not None:
            cmds.progressBar(self.bar, edit=True, step=1, status=stts)
        else:
            cmds.progressBar(self.bar, edit=True, step=1)
        self.count += 1

    def end(self):
        cmds.progressBar(self.bar, edit=True, endProgress=True)

    def isCancelled(self):
        return cmds.progressBar(self.bar, q=True, isCancelled=True)

    def __enter__(self):
        return self
    def __exit__(self, type, value, traceback):
        self.end()




#プリファレンスのスムーズタイプをコマンドで使用可能なサブディビジョンタイプに変換するためのテーブル
_gblSmoothDrawTypeToSubdivisionType = (0, 2, 3)

#スムーズメッシュプレビューをポリゴンに変換
# isRendering : True    レンダリング用のスムーズレベルを使用する
#             : False   プレビュー用のスムーズレベルを使用する
#
# return       : ([smooths], [histories])   追加されたスムースノードとコンストラクションヒストリをタプルで返す
#
def toPolySmooth(shape, isRendering):
    if cmds.getAttr(shape + ".useGlobalSmoothDrawType"):
        smoothDrawType = cmds.optionVar(q="globalPolyDefaultSmoothDrawType")
        subdivisionType = _gblSmoothDrawTypeToSubdivisionType[smoothDrawType]
    else:
        smoothDrawType = cmds.getAttr(shape + ".smoothDrawType")
        subdivisionType = smoothDrawType

    if not isRendering or cmds.getAttr(shape + ".useSmoothPreviewForRender"):
        smoothLevel = cmds.getAttr(shape + ".smoothLevel")
    else:
        smoothLevel = cmds.getAttr(shape + ".renderSmoothLevel")

    #Maya Catmull-Clark
    continuity = cmds.getAttr(shape + ".continuity")        #連続性
    boundaryRule = cmds.getAttr(shape + ".boundaryRule")    #境界の基準
    smoothUVs = cmds.getAttr(shape + ".smoothUVs")          #UVのスムーズ
    propagateEdgeHardness = cmds.getAttr(shape + ".propagateEdgeHardness")    #エッジ硬度の伝播
    keepMapBorders = cmds.getAttr(shape + ".keepMapBorders")#マップ境界
    keepBorder = cmds.getAttr(shape + ".keepBorder")        #ジオメトリ境界
    keepHardEdge = cmds.getAttr(shape + ".keepHardEdge")    #ハードエッジ

    #OpenSubdiv
    osdVertBoundary = cmds.getAttr(shape + ".osdVertBoundary")  #頂点の境界
    osdFvarBoundary = cmds.getAttr(shape + ".osdFvarBoundary")  #UV境界のスムージング
    osdFvarPropagateCorners = cmds.getAttr(shape + ".osdFvarPropagateCorners")  #UVコーナーを伝播
    osdSmoothTriangles = cmds.getAttr(shape + ".osdSmoothTriangles")    #スムーズ三角形
    osdCreaseMethod = cmds.getAttr(shape + ".osdCreaseMethod")  #折り目メソッド

    #constructionHistory = cmds.constructionHistory(q=True, tgl=True)
    constructionHistory = True

    org_histories = cmds.listHistory(shape)

    cmds.setAttr(shape + ".displaySmoothMesh", 0)
    #cmds.setAttr(shape + ".smoothLevel", 2)

    smooths = cmds.polySmooth(shape, method = 0
                        , sdt = subdivisionType
                        , dv = smoothLevel
                        , c = continuity
                        , bnr = boundaryRule
                        , suv = smoothUVs
                        , peh = propagateEdgeHardness
                        , kmb = keepMapBorders
                        , kb = keepBorder
                        , khe = keepHardEdge
                        , ovb = osdVertBoundary
                        , ofb = osdFvarBoundary
                        , ofc = osdFvarPropagateCorners
                        , ost = osdSmoothTriangles
                        , ocr = osdCreaseMethod
                        #, ksb = True
                        #, kt = False
                        , ch = constructionHistory
                        ) or []
    #cmds.polySelect(shape)  #スムース化した際にエッジが選択されてしまうのでキャンセルする

    # コンストラクションヒストリが無いshapeの場合、自動でpolySurfaceShapeが追加されるので、増加したヒストリを調べる
    added_histories = [hist for hist in cmds.listHistory(shape) if hist not in org_histories and not hist in smooths]
#    print 'org = {}'.format(org_histories)
#    print 'add = {}'.format(added_histories)
    return (smooths, added_histories)


#スムーズメッシュプレビューをポリゴンに一括変換
def toPolySmooths(shapes, isRendering):
    #ビューポートレンダー時（現在未使用）
    if not isRendering:
        smooths_and_hists = []
        for shape in shapes:
            result = toPolySmooth(shape, isRendering)
            if result[0] or result[1]:
                smooths_and_hists.append(result)


    #バッチレンダー時（プログレス表示なし）
    if cmds.about(q=True, batch=True):
        smooths_and_hists = []
        for i, shape in enumerate(shapes):
            plogger.info('{:3} : '.format(i) + shape)   #shape名にAscii以外が含まれているとformat()でエラーになるので外に出す
            result = toPolySmooth(shape, isRendering)
            if result[0] or result[1]:
                smooths_and_hists.append(result)
        return smooths_and_hists


    #レンダービュー
    #キャンセルフラグをクリアするためにプログレスバーを空回しする
    with MainProgress(1, True, '') as progress:
        progress.step()

    smooths_and_hists = []
    with MainProgress(len(shapes), True) as progress:
        for i, shape in enumerate(shapes):
            if progress.isCancelled():
                plogger.info('cancelled !!')
                break
            plogger.info('{:3} : '.format(i) + shape)
            progress.step('to polygons : ' + shape)
            result = toPolySmooth(shape, isRendering)
            if result[0] or result[1]:
                smooths_and_hists.append(result)

    #helpLine update
    with MainProgress(1, False, '') as progress:
        progress.step()

    return smooths_and_hists


#スムースメッシュプレビューが有効なシェイプを取得する
def getSmoothPreviewShapes(isRendering, isCurrentLayerOnly):

    #可視状態の判定（transform, meshには同じチャンネルがあるので両対応）
    vis_judge = VisibleJudgment()
    _isVisible = vis_judge.isVisible

    #可視状態の判定（transform, meshには同じチャンネルがあるので両対応）
#    def _isVisible(node):
#        if not node or not cmds.getAttr(node + '.visibility') or cmds.getAttr(node + '.template'):
#            return False
#        if cmds.getAttr(node + '.overrideEnabled'):
#            if not cmds.getAttr(node + '.overrideVisibility') or cmds.getAttr(node + '.overrideDisplayType') == 1:    # 1 == Template
#                return False
#        return True

    #親を通して可視状態の判定
    def _isVisibleTree(nodepath):
        paths = nodepath.split('|')
        for i in range(len(paths)-1):
            node = '|'.join(paths[:i+2])
            if not _isVisible(node):
                return False
        return True

    #スムーズ変換の必要があるメッシュかの判定
    def _isSmoothTargetMesh(mesh):
        if cmds.getAttr(mesh + ".displaySmoothMesh") and not cmds.getAttr(mesh + ".intermediateObject") and cmds.getAttr(mesh + ".allowTopologyMod"):
            if not isRendering or cmds.getAttr(mesh + ".useSmoothPreviewForRender"):
                return cmds.getAttr(mesh + ".smoothLevel") > 0
            else:
                return cmds.getAttr(mesh + ".renderSmoothLevel") > 0
        return False

    #メッシュがインスタンス化されている場合に元のメッシュを取得する
    _orig_meshes = cmds.ls(long=True, exactType=('mesh'), noIntermediate=True)
    def _getFirstMeshInstance(mesh):
        instances = cmds.ls(mesh, long=True, allPaths=True)
        if len(instances) > 1:  #親が複数ある場合はインスタンス化されている
            for inst in instances:
                if inst in _orig_meshes:
                    return inst
        return instances[0]


    #可視状態のトランスフォームを末端に向けてたどりスムーズ対象メッシュを収集する
    def _recursChildMeshes(trans):
        #
        def _recursChildSub(trans, shapes):
            for node in cmds.listRelatives(trans, fullPath=True) or []:
                node_type = cmds.nodeType(node)
                if node_type == 'transform' and _isVisible(node):
                    _recursChildSub(node, shapes)
                elif node_type == 'mesh' and _isSmoothTargetMesh(node) and _isVisible(node):
                    shapes.append(_getFirstMeshInstance(node))
        shapes = []
        _recursChildSub(trans, shapes)
        return shapes

    #DAGパスソート用の比較関数
    def _dagPathCmp(x_, y_):
        xx = x_.split('|')
        yy = y_.split('|')
        xx_n = len(xx)
        yy_n = len(yy)
        if xx_n != yy_n:
            return -1 if xx_n < yy_n else 1
        for x, y in zip(xx, yy):
            x_n = len(x)
            y_n = len(y)
            if x_n != y_n:
                return -1 if x_n < y_n else 1
            r = cmp(x, y)
            if r != 0:
                return r
        return 0

    #ツリーを作成し上位パスのみ収集する
    #例
    # ['|pCube1|pCubeShape1', '|pCube2', '|pCube1']
    # => ['|pCube2', '|pCube1']     # '|pCube1|pCubeShape1'は'|pCube1'の下位パスなので省略される
    class TreeNode(object):
        def __init__(self, name, parent):
            self.name = name
            self.parent = parent
            self.childs = {}

        def _recursParent(self, paths):
            if self.parent != None:
                self.parent._recursParent(paths)
            paths.append(self.name)

        def fullName(self):
            paths = []
            self._recursParent(paths)
            return '|'.join(paths)

        @classmethod
        def MakeUpperTree(cls, paths):
            cls.root = TreeNode('', None)
            for path in paths:
                node = cls.root
                _paths = path.split('|')[1:]
                _paths_n = len(_paths)
                for i, name in enumerate(_paths, 1):
                    if name in node.childs:
                        if not node.childs[name].childs:    #既存のパスの方が短いので中断
                            break;
                        if i == _paths_n:                   #今のパスの方が短いので置き換え
                            node.childs[name].childs = {}
                            break
                    else:
                        node.childs[name] = TreeNode(name, node)
                    node = node.childs[name]

        @classmethod
        def _recursChild(cls, node, nodes):
            for child in node.childs.values():
                fullname = child.fullName()
                if _isVisible(fullname):
                    if child.childs:
                        cls._recursChild(child, nodes)
                    else:
                        nodes.append(fullname)

        @classmethod
        def getVisibleNodes(cls):
            nodes = []
            cls._recursChild(cls.root, nodes)
            return nodes;

        @classmethod
        def clean(cls):
            if hasattr(cls, 'root'):
                del cls.root


    if isCurrentLayerOnly:
        #PencilrenderSettingノードを用いた場合はレンダーレイヤ毎にレンダーコールバックが呼び出されるので、カレントレンダーレイヤのみを対象とするだけで良い
        plogger.info('current layer only')
        layer = cmds.editRenderLayerGlobals(q=True, currentRenderLayer=True)
        nodes = cmds.editRenderLayerMembers(layer, q=True, fullNames=True) or []
        plogger.info(layer)
    else:
        #レンダー設定のプリレンダーMELから使用される場合は、レンダリング可能なレンダーレイヤをまとめて処理する
        #※カレントレイヤ―がレンダリング可能とは限らない
        plogger.info('all layers')
        nodes = []
        for layer in cmds.listConnections('renderLayerManager.renderLayerId'):
            if cmds.getAttr(layer + '.renderable'):
                plogger.info(layer)
                nodes += cmds.editRenderLayerMembers(layer, q=True, fullNames=True)
        nodes = set(nodes)

    transforms = []
    meshes = []
    for node in nodes:
        try:
            node_type = cmds.nodeType(node)
            if node_type == 'transform':
                transforms.append(node)
            elif node_type == 'mesh':
                if _isSmoothTargetMesh(node):
                    meshes.append(node)
        except:
            plogger.info('skip node = ' + node)
            pass

    #インスタンスメッシュを展開してから可視判定をする
    if meshes:
        expan_meshes = [p + '|' + mesh.split('|')[-1] for mesh in meshes for p in cmds.listRelatives(mesh, fullPath=True, allParents=True)]
        meshes = [_getFirstMeshInstance(mesh) for mesh in set(expan_meshes) if _isVisibleTree(mesh)]
        meshes = list(set(meshes))


    #トランスフォームのツリーをまとめる
    TreeNode.MakeUpperTree(transforms)
    transforms = TreeNode.getVisibleNodes()
    TreeNode.clean()

    #トランスフォームの場合は再帰的に子供を含む
    for trans in transforms:
        meshes += _recursChildMeshes(trans)

    return sorted(set(meshes), key=functools.cmp_to_key(_dagPathCmp))


#シェイプのスムースメッシュプレビューフラグ値を取得する
def getSmoothPreviewFlags(shapes):
    flags = [ cmds.getAttr(shape + ".displaySmoothMesh") for shape in shapes ]
    return flags


#シェイプのソースコネクションを取得する
def getSourceConnections(shapes):
    connections = [ cmds.listConnections(shape, source=True, destination=False, plugs=True, connections=True) or [] for shape in shapes ]
    return connections;


#スムースフェースの削除を行い、シェイプのスムースメッシュプレビューフラグを再セットする
def restoreSmoothPreviewFlagsAndRemoveSmooths(shapes, flags, connections, smooths_and_hists):

    #スムースオブジェクトの削除
    smooths = [ smooth for smooths_, hists_ in smooths_and_hists for smooth in smooths_ ]  # flatten smooths
    if smooths:
        with MainProgress(len(smooths), False, 'restoreSmoothPreviewFlagsAndRemoveSmooths') as progress:
            plogger.info('remove smooths')
            for i, smooth in enumerate(smooths):
                plogger.info('{:3} : '.format(i) + smooth)
                progress.step('remove smooths : ' + smooth)
                cmds.delete(smooth)

    #スムースフラグ復帰
    for shape, flag in zip(shapes, flags):
        cmds.setAttr(shape + ".displaySmoothMesh", flag)
        cmds.polySelect(shape)  #スムースを除去した際にエッジが選択されてしまうのでキャンセルする

    #新たにヒストリが追加されてしまったものは削除する
    remove_shape_hists = [ (shape_, smooth_and_hist_[1]) for shape_, smooth_and_hist_ in zip(shapes, smooths_and_hists) if len(smooth_and_hist_[1]) > 0 ]
    if remove_shape_hists:
        with MainProgress(len(remove_shape_hists), False, 'restoreSmoothPreviewFlagsAndRemoveSmooths') as progress:
            plogger.info('remove construction history')
            for i, shape_hist in enumerate(remove_shape_hists):
                try:
                    shape = shape_hist[0]
                    hists = shape_hist[1]
                    plogger.info('{:3} : '.format(i) + shape)
                    progress.step('remove histories : ' + shape)
                    #リファレンスかどうかで除去方法を変える
                    if cmds.referenceQuery(shape, isNodeReferenced=True):
                        for hist in hists:
                            try:
                                plogger.info(' ^^^^^ ' + hist)
                                cmds.delete(hist)
                            except ValueError:
                                pass
                    else:
                        cmds.delete(shape, constructionHistory=True)
                except ValueError:
                    pass

    #ソースコネクションの復帰（リファレンスシーンの場合にコネクションが切断されてしまうケースがあるので復帰する）
    for conns in connections:
        for i in range(0, len(conns), 2):
            try:
                dst = conns[i]
                src = conns[i+1]
                if not cmds.isConnected(src, dst):
                    cmds.connectAttr(src, dst)
            except RuntimeError:
                pass


#セレクトエッジに接続されているセットを取得する
def getSetsConnectedToSelectEdges():
    return list(set(
            [_set for edge in cmds.ls(long=True, type='PencilSelectedEdge')
                for _set in (cmds.listConnections(edge + '.message', destination=True, type='objectSet') or []) if cmds.sets(_set, q=True, size=True) > 0]
            ))

#セットのメンバーを取得する
def getSetsMembers(_sets):
    return [ cmds.sets(_set, q=True) for _set in _sets ]

#セットメンバーを削除して復帰する
def restoreSetsMembers(_sets, _sets_members):
    for _set, _member in zip(_sets, _sets_members):
        members = cmds.sets(_set, q=True)
        cmds.sets(members, e=True, rm=_set)
        for i in range(3):
            cmds.sets(_member, e=True, add=_set)    #2回書き込まないと反映されないことがあるためループで何回か試す（リファレンスシーンで確認）
            members = cmds.sets(_set, q=True)
            if members:
                break;
            else:
                plogger.info('try {} : {} : Failed to add SET member'.format(i, _set))

#セレクトエッジを列挙しソース側のリンクを切断する
def disconnectSelectedEdge():
    seledges = [(pedge
            , cmds.connectionInfo(pedge + '.mesh', sourceFromDestination=True)
            , cmds.connectionInfo(pedge + '.meshReference', sourceFromDestination=True))
            for pedge in cmds.ls(long=True, type='PencilSelectedEdge')]
    for e in seledges:
        if e[1]:
            cmds.disconnectAttr(e[1], e[0] + '.mesh')
        if e[2]:
            cmds.disconnectAttr(e[2], e[0] + '.meshReference')
    return seledges

#セレクトエッジのソース側のリンクを復帰する
def restoreSelectedEdge(seledges):    
    for e in seledges:
        if e[1]:
            cmds.connectAttr(e[1], e[0] + '.mesh')
        if e[2]:
            cmds.connectAttr(e[2], e[0] + '.meshReference')

#セレクトエッジが接続されたままスムーズ変換を行うと処理速度が低下するため、一旦切断後に再接続する
class DisconnectSelectedEdge(object):
    def __enter__(self):
        self.seledges = disconnectSelectedEdge()
        return self
    def __exit__(self, type, value, traceback):
        restoreSelectedEdge(self.seledges)




_invalidId = -1
_dictKey = 0
_dict = {}

def Reset():
    global _dictKey, _dict
    _dictKey = 0
    _dict = {}

class Restore(object):
    def __init__(self, shapes, flags, connections, smooths_and_hists, sets, sets_members, selections, current_ctx):
        self.shapes = shapes
        self.flags = flags
        self.connections = connections
        self.smooths_and_hists = smooths_and_hists
        self.sets = sets
        self.setsMembers = sets_members
        self.selections = selections
        self.currentCtx = current_ctx

    def __del__(self):
        with DisconnectSelectedEdge():
            restoreSmoothPreviewFlagsAndRemoveSmooths(self.shapes, self.flags, self.connections, self.smooths_and_hists)
        with MainProgress(2, False, 'restoreSetsMembers') as progress:
            restoreSetsMembers(self.sets, self.setsMembers)
            progress.step()
            cmds.select(self.selections)
            cmds.setToolTo(self.currentCtx)
            progress.step()
        plogger.info('restored shapes cnt = {}'.format(len(self.shapes)))
        plogger.info('restored sets cnt = {}'.format(len(self.sets)))


def SmoothBefore(isRendering=True, isCurrentLayerOnly=True):
    """
    isRendering : True  レンダリング用
                : Flase View Port用
    isCurrentLayerOnly : True   カレントレンダーレイヤのみ変換対象とする
                       : False  全レンダーレイヤを変換対象とする
    """
    try:
        global _dictKey, _dict

        _current_ctx = cmds.currentCtx()
        # 選択ツール以外の場合に下記の不具合があるので、事前に選択ツールに切り替える
        # 1.レンダーレイヤのメンバーにDAG内に存在しないマニュピレータオブジェクトが追加される
        # 2.スムースポリゴン化の際にフリーズすることがある
        cmds.setToolTo('selectSuperContext');

        plogger.info(cmds.about(iv=True))
        _shapes = getSmoothPreviewShapes(isRendering, isCurrentLayerOnly)
        if len(_shapes) == 0:
            cmds.setToolTo(_current_ctx)
            plogger.info('No target node')
            return _invalidId

        plogger.info('smooth shapes cnt = {}'.format(len(_shapes)))
        _selections = cmds.ls(selection=True)
        _flags = getSmoothPreviewFlags(_shapes)
        _connections = getSourceConnections(_shapes)
        _sets = getSetsConnectedToSelectEdges()
        _sets_members = getSetsMembers(_sets)
        plogger.info('selected edge sets cnt = {}'.format(len(_sets)))
        with DisconnectSelectedEdge():
            _smooths_and_hists = toPolySmooths(_shapes, isRendering)
        cmds.polySelect(_shapes)  #スムース化した際にエッジが選択されてしまうのでキャンセルする
        if len(_dict) == 0:
            _dictKey = 0
        else:
            while (_dictKey in _dict):
                _dictKey += 1
        _dict[_dictKey] = Restore(_shapes, _flags, _connections, _smooths_and_hists, _sets, _sets_members, _selections, _current_ctx)
        plogger.info('restore id = {}'.format(_dictKey))
        return _dictKey
    except Exception as e:
        plogger.info('error !!')
        plogger.info(e)
        raise

def SmoothAfter(id):
    plogger.info('restore id = {}'.format(id))
    try:
        del _dict[id]
    except KeyError:
        return False
    return True




#
# Maya Software Render用 (PencilRenderSettingが有効な場合）
# ※レンダーレイヤが使用されている場合、レイヤ毎に呼び出されます
def Smooth_RenderBefore():
    return SmoothBefore(True, True)

def Smooth_RenderAfter(id):
    SmoothAfter(id)

def Smooth_RenderInterrupted(id):
    plogger.info('interrupted !!!!!')
    SmoothAfter(id)

#
# レンダー設定のプリレンダーMEL用
#
def Smooth_PreRender():
    is_batch = cmds.about(q=True, batch=True)
    return SmoothBefore(True, not is_batch)

def Smooth_PostRender(id):
    SmoothAfter(id)


#
# View Port用
# ※未使用
def Smooth_ViewPortBefore():
    #cmds.refresh(suspend=True)
    return SmoothBefore(False)

def Smooth_ViewPortAfter(id):
    SmoothAfter(id)
    #cmds.refresh(suspend=False)
